import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.font_manager as fm # 新增导入

# 设置中文字体，以确保图表中的中文能正常显示
# 设置中文字体为宋体
# 尝试查找系统中的宋体，如果找不到则使用默认的 sans-serif 字体
try:
    # 查找宋体（SimSun）
    font_path_simsun = fm.findfont(fm.FontProperties(family='SimSun'))
    plt.rcParams['font.sans-serif'] = ['SimSun', 'Times New Roman'] # 优先使用宋体，备用新罗马
except:
    print("未找到宋体（SimSun），将使用默认字体。")
    plt.rcParams['font.sans-serif'] = ['Times New Roman', 'DejaVu Sans'] # 如果宋体找不到，确保英文和数字是新罗马

# 设置英文和数字字体为新罗马
plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['font.serif'] = ['Times New Roman']
plt.rcParams['font.monospace'] = ['Times New Roman']

plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方块的问题

def find_fft_data_path(root_path, flow_rate):
    """
    根据流速查找对应的 fft_data_X.csv 文件路径。
    """
    # 构建预期的文件夹名称，例如 1000f0.1ms, 1000f3.0ms
    # 对于整数流速，例如 3.0，格式化为 "3.0"
    flow_rate_str = f"{flow_rate:.1f}"
    folder_name = f"1000f{flow_rate_str}ms"

    processed_folders = [d for d in os.listdir(root_path) if d.endswith('_processed')]
    if not processed_folders:
        print(f"在 {root_path} 中未找到 'processed' 后缀的文件夹")
        return None

    # 假设我们使用第一个带 'processed' 后缀的文件夹
    search_path_base = os.path.join(root_path, processed_folders[2])
    
    flow_folder_path = os.path.join(search_path_base, folder_name)

    if os.path.isdir(flow_folder_path):
        # 查找下一级目录，通常是 'C001H001S0001'
        subfolders = [d for d in os.listdir(flow_folder_path) if os.path.isdir(os.path.join(flow_folder_path, d))]
        if subfolders:
            # 假设数据在第一个子文件夹中
            data_folder_path = os.path.join(flow_folder_path, subfolders[0])
            fft_file_path = os.path.join(data_folder_path, 'fft_data_X.csv')
            if os.path.exists(fft_file_path):
                return fft_file_path

    print(f"在 {search_path_base} 中未找到流速为 {flow_rate} m/s 的数据文件路径（预期文件夹：{folder_name}）。")
    return None


def plot_spectrum(data_groups, root_path, output_folder):
    """
    根据提供的数据分组绘制频谱图。

    :param data_groups: 一个字典，键是图表标题，值是流速列表。
    :param root_path: 数据存储的根目录。
    :param output_folder: 保存生成图表的文件夹。
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    for i, (title, flow_rates) in enumerate(data_groups.items()):
        plt.figure(figsize=(12, 8))
        
        for flow_rate in flow_rates:
            file_path = find_fft_data_path(root_path, flow_rate)
            
            if file_path:
                # 打印当前子图对应的数据文件夹路径
                print(f"正在处理图表: {title}, 流速: {flow_rate} m/s, 数据来源文件夹: {os.path.dirname(file_path)}")
                try:
                    # 读取CSV文件，跳过可能存在的表头
                    df = pd.read_csv(file_path, header=None, names=['frequency', 'amplitude'], skiprows=1)
                    
                    # 将列转换为数值类型，无法转换的将变为NaN
                    df['frequency'] = pd.to_numeric(df['frequency'], errors='coerce')
                    df['amplitude'] = pd.to_numeric(df['amplitude'], errors='coerce')
                    
                    # 删除包含NaN的行
                    df.dropna(inplace=True)

                    # 筛选频率小于15Hz的数据
                    df_filtered = df[(df['frequency'] > 0.2) & (df['frequency'] < 10)]
                    
                    # 绘制频谱曲线
                    plt.plot(df_filtered['frequency'], df_filtered['amplitude'], label=f'{flow_rate} m/s', linewidth=0.5)
                except Exception as e:
                    print(f"处理文件 {file_path} 时出错: {e}")

        # 明确设置标题、轴标签和图例的字体
        plt.title(title, fontsize=16, fontproperties=fm.FontProperties(family=['SimSun', 'Times New Roman']))
        plt.xlabel('频率 (Hz)', fontsize=18, fontproperties=fm.FontProperties(family=['SimSun', 'Times New Roman']))
        plt.ylabel('幅值', fontsize=18, fontproperties=fm.FontProperties(family=['SimSun', 'Times New Roman']))
        plt.legend(fontsize=21, prop=fm.FontProperties(family=['SimSun', 'Times New Roman']))
        plt.grid(True)
        plt.tick_params(axis='both', which='major', labelsize=14)
        
        # 确保刻度标签的字体也设置为新罗马
        for label in plt.gca().get_xticklabels() + plt.gca().get_yticklabels():
            label.set_fontname('Times New Roman')
        
        # 保存图表
        output_path = os.path.join(output_folder, f'spectrum_plot_{i+1}.png')
        plt.savefig(output_path, dpi=300)
        plt.close()
        print(f"图表已保存至: {output_path}")


if __name__ == '__main__':
    # 定义数据根目录
    ROOT_DATA_PATH = r'J:\lqb\huanxingranliao\25CHJ02-cu'
    
    # 定义输出文件夹
    OUTPUT_FOLDER = 'spectrum_analysis_results'

    # 定义图表分组和对应的流速
    # 注意：这里的流速值是浮点数
    plot_groups = {
        '频谱曲线图 (0.1-1.5 m per s)': [0.1, 0.6, 0.9, 1.5],
        '频谱曲线图 (1.8-3.2 m per s)': [1.8, 2.4, 2.8, 3.2],
        '频谱曲线图 (3.5-5.0 m per s)': [3.5, 4.1, 4.7, 5.0]
    }

    # 执行绘图
    plot_spectrum(plot_groups, ROOT_DATA_PATH, OUTPUT_FOLDER)
